#!/bin/env python
# -*- coding: utf-8 -*-

#  This file is part of GLExDrv_DSO8060.
#
#  GLExDrv_DSO8060 is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  GLExDrv_DSO8060 is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with GLExDrv_DSO8060.  If not, see <http://www.gnu.org/licenses/>.


# Simple driver for the USB interface on the Hantek DSO-8060 portable
# multifunction instrument. Currently only controls the signal generator.
# (C) 2012 Greg Courville <greg_courville@greglabs.com>

import usb
from math import pi

def to_uint16_be(x):
	return x & 0xFF, (x & 0xFF00) >> 8

class USBInstrument(object):
	class NoDevicesFound(StandardError):
		pass
	class RemoteError(StandardError):
		pass
	@classmethod
	def all_devices(cls):
		for bus in usb.busses():
			for dev in bus.devices:
				if dev.idVendor == cls.VENDOR_ID \
				 and dev.idProduct == cls.PRODUCT_ID:
					h = dev.open()
					serial = h.getString(dev.iSerialNumber, 255)
					yield dev, serial
	@classmethod
	def first_device(cls):
		try:
			return cls.all_devices().next()
		except StopIteration:
			raise cls.NoDevicesFound()
	def __init__(self, usb_dev = None):
		if usb_dev == None:
			usb_dev = self.first_device()[0]
		self.usb_dev = usb_dev
		self.usb_handle = usb_dev.open()
	def close(self):
		if hasattr(self, "usb_handle"):
			self.usb_handle.releaseInterface()
	def __del__(self):
		self.close()

class DSO8060USB(USBInstrument):
	VENDOR_ID = 0x4348
	PRODUCT_ID = 0x5538
	def __init__(self, usb_dev = None):
		USBInstrument.__init__(self, usb_dev)
		self.usb_handle.claimInterface(0)
	def _check_command_result(self):
		r = self.usb_handle.bulkRead(0x82, 64)
		if len(r) == 1:
			if r[0] == 132:
				return
		raise self.RemoteError()		
	def _send_command(self, a, b_lo, b_hi, c_lo, c_hi):
		self.usb_handle.bulkWrite(
			0x02,
			chr(a) + "\x00\x00\x00" + chr(b_lo) + chr(b_hi) + "\x00\x00"
			 + chr(c_lo) + chr(c_hi) + "\x00\x00"
			)
	def setmode_dds(self):
		self._send_command(0xA0, 0x3D, 0, 0x02, 0)
	def setmode_dmm(self):
		self._send_command(0xA0, 0x3D, 0, 0x01, 0)
	def setmode_scope(self):
		self._send_command(0xA0, 0x3D, 0, 0x00, 0)
	def dds_set_clkdiv(self, n):
		lo, hi = to_uint16_be(n-1)
		self._send_command(0xA0, 0x3C, 0, lo, hi)
	def _dds_set_length(self, n):
		lo, hi = to_uint16_be(n * 2)
		self._send_command(0xA2, lo, hi, 0x09, 0)
		self.dds_length = n
	def dds_write_raw(self, pts):
		self._dds_set_length(len(pts))
		data = []
		for pt in pts:
			lo, hi = to_uint16_be(pt)
			data += [lo,hi]
		self.usb_handle.bulkWrite(0x02, data)
		self._check_command_result()
	def dds_write_volts(self, pts):
		self.dds_write_raw([int(2048 * (1. - pt/3.5)) for pt in pts])
	def dds_set_poweron(self, on):
		self._send_command(0xA0, 0x3B, 0, on and 1 or 0, 0)
	def dds_set_sync_out(self, on):
		self._send_command(0xA0, 0x3A, 0, on and 1 or 0, 0)
	def dds_fill_from_fn(self, fn, num_pts=2048, start=0., stop=2.*pi, phase=0, scale=1.0, offset=0.):
		step = float(stop-start) / num_pts
		self.dds_write_volts(
			[scale * fn(i * step + start + phase) + offset for i in range(num_pts)]
			)

